/**
 *  Copyright (C) 2011
 *  University of Rochester Department of Computer Science
 *    and
 *  Lehigh University Department of Computer Science and Engineering
 *
 * License: Modified BSD
 *          Please see the file LICENSE.RSTM for licensing information
 */

#ifndef BMHARNESS_HPP__
#define BMHARNESS_HPP__

#include <cstdlib>
#include <iostream>
#include <signal.h>
#include <pthread.h>
#include <rstm.hpp>
#include "bmconfig.hpp"
#include "../include/abstract_timing.hpp"
#include "../include/abstract_cpu.hpp"

#ifdef STM_HAS_AOU
#include "../include/ptlcalls.h"
#else
#define ptlcall_switch_to_sim()    ((void)0)
#define ptlcall_switch_to_native() ((void)0)
#endif

using std::string;
using std::cout;

Config::Config() :
    bmname(""),
    duration(1),
    execute(0),
    threads(1),
    nops_after_tx(0),
    elements(256),
    lookpct(34),
    inspct(66),
    sets(1),
    ops(1),
    switch_to_sim(0),
    time(0),
    running(true),
    txcount(0)
{
}

Config CFG TM_ALIGN(64);

namespace
{

  /**
   * Print benchmark configuration output
   */
  void dump_csv()
  {
      // csv output
      std::cout << "csv"
                << ", ALG=" << TM_GET_ALGNAME()
                << ", B=" << CFG.bmname     << ", R=" << CFG.lookpct
                << ", d=" << CFG.duration   << ", p=" << CFG.threads
                << ", X=" << CFG.execute    << ", m=" << CFG.elements
                << ", S=" << CFG.sets       << ", O=" << CFG.ops
                << ", txns=" << CFG.txcount << ", time=" << CFG.time
                << ", throughput="
                << (1000000000LL * CFG.txcount) / (CFG.time)
                << std::endl;
  }

  /**
   *  Print usage
   */
  void usage()
  {
      std::cerr << "Usage: CounterBench -C <stm algorithm> [flags]\n";
      std::cerr << "    -d: number of seconds to time (default 1)\n";
      std::cerr << "    -X: execute fixed tx count, not for a duration\n";
      std::cerr << "    -p: number of threads (default 1)\n";
      std::cerr << "    -N: nops between transactions (default 0)\n";
      std::cerr << "    -R: % lookup txns (remainder split ins/rmv)\n";
      std::cerr << "    -m: range of keys in data set\n";
      std::cerr << "    -B: name of benchmark\n";
      std::cerr << "    -S: number of sets to build (default 1)\n";
      std::cerr << "    -O: operations per transaction (default 1)\n";
      std::cerr << "    -s: switch to the simulator\n";
      std::cerr << "    -h: print help (this message)\n\n";
  }

/**
 *  Parse command line arguments
 */
void
parseargs(int argc, char** argv)
{
    // parse the command-line options
    int opt;
    while ((opt = getopt(argc, argv, "N:d:p:hX:B:m:R:S:O:s")) != -1) {
        switch(opt) {
          case 'd': CFG.duration      = strtol(optarg, NULL, 10); break;
          case 'p': CFG.threads       = strtol(optarg, NULL, 10); break;
          case 'N': CFG.nops_after_tx = strtol(optarg, NULL, 10); break;
          case 'X': CFG.execute       = strtol(optarg, NULL, 10); break;
          case 'B': CFG.bmname        = std::string(optarg); break;
          case 'm': CFG.elements      = strtol(optarg, NULL, 10); break;
          case 'S': CFG.sets          = strtol(optarg, NULL, 10); break;
          case 'O': CFG.ops           = strtol(optarg, NULL, 10); break;
          case 's': CFG.switch_to_sim = 1;                        break;
          case 'R':
            CFG.lookpct = strtol(optarg, NULL, 10);
            CFG.inspct = (100 - CFG.lookpct)/2 + strtol(optarg, NULL, 10);
            break;
          case 'h':
            usage();
        }
    }
}

/**
 *  Run some nops between transactions, to simulate some time being spent on
 *  computation
 */
void nontxnwork()
{
    if (CFG.nops_after_tx)
        for (uint32_t i = 0; i < CFG.nops_after_tx; i++)
            nop();
}

/*** Signal handler to end a test */
extern "C" void catch_SIGALRM(int) {
    CFG.running = false;
}

/**
 *  Support a few lightweight barriers
 */
void
barrier(uint32_t which)
{
    static volatile uint32_t barriers[16] = {0};
    CFENCE;
    fai32(&barriers[which]);
    while (barriers[which] != CFG.threads) { }
    CFENCE;
}

/**
 * Move the calling thread to the specified CPU.
 */
void move_to_core(int i) {
    cpu_set_t cpu;
    CPU_ZERO(&cpu);
    CPU_SET(i, &cpu);
    //sched_setaffinity(0, sizeof(cpu), &cpu);
    pthread_setaffinity_np(pthread_self(), sizeof(cpu), &cpu);
    //printf("I want to run on %i and ended up on core %i\n", i, sched_getcpu());
#if (0)
    printf("I want to run on %i and ended up on core %i\n", i, omp_get_thread_num());
    
    // Getting the actual masks back
    pthread_getaffinity_np(pthread_self(), sizeof(cpu), &cpu);
    for (int i = 0; i < 32; i++)
        if (CPU_ISSET(i, &cpu))
            printf ("Pthread Affinity: %i\n", i);

    sched_getaffinity(0, sizeof(cpu), &cpu);
    for (int i = 0; i < 32; i++)
        if (CPU_ISSET(i, &cpu))
            printf ("Sched Affinity: %i\n", i);
    printf("Sched Get CPU: %i\n", sched_getcpu());
#endif
}

/*** Run a timed or fixed-count experiment */
void
run(uintptr_t id)
{
    // create a transactional context (repeat calls from thread 0 are OK)
    TM_THREAD_INIT();

    // wait until all threads created, then set alarm and read timer
    barrier(0);
    if (id == 0) {
        if (!CFG.execute) {
            signal(SIGALRM, catch_SIGALRM);
            alarm(CFG.duration);
        }
        CFG.time = getElapsedTime();
    }

    // wait until read of start timer finishes, then start transactions
    barrier(1);
#if defined STM_CHECKPOINT_ASM
    volatile uint32_t count = 0;
#else
    uint32_t count = 0;
#endif
    uint32_t seed = id; // not everyone needs a seed, but we have to support it
    if (!CFG.execute) {
        // run txns until alarm fires
        while (CFG.running) {
            bench_test(id, &seed);
            ++count;
            nontxnwork(); // some nontx work between txns?
        }
    }
    else {
        // run fixed number of txns
#if defined STM_CHECKPOINT_ASM
        for (volatile uint32_t e = 0; e < CFG.execute; e++)
#else
        for (uint32_t e = 0; e < CFG.execute; e++)
#endif
        {
            bench_test(id, &seed);
            ++count;
            nontxnwork(); // some nontx work between txns?
        }
    }

    // wait until all txns finish, then get time
    barrier(2);
    if (id == 0)
        CFG.time = getElapsedTime() - CFG.time;

    // add this thread's count to an accumulator
    faa32(&CFG.txcount, count);
}

/**
 *  pthread wrapper for running the experiments
 *
 *  NB: noinline prevents this from getting inlined into main (and prevents
 *      run from being inlined there as well. This eliminates an
 *      _ITM_initializeProcess ordering problem if there's a transaction
 *      lexically scoped inside of main.
 */
NOINLINE
void*
run_wrapper(void* i)
{
    move_to_core((int)(long)i);

    if (CFG.switch_to_sim) {
        barrier(14);
        if (!i)
            ptlcall_switch_to_sim();
        barrier(15);
    }

    run((uintptr_t)i);
    TM_THREAD_SHUTDOWN();
    return NULL;
}
} // namespace

/**
 *  Main routine: parse args, set up the TM system, prep the benchmark, run
 *  the experiments, verify results, print results, and shut down the system
 */
int main(int argc, char** argv) {
    parseargs(argc, argv);
    bench_reparse();
    TM_SYS_INIT();
    TM_THREAD_INIT();
    bench_init();

    void* args[256];
    pthread_t tid[256];

    // set up configuration structs for the threads we'll create
    pthread_attr_t attr;
    pthread_attr_init(&attr);
    pthread_attr_setscope(&attr, PTHREAD_SCOPE_SYSTEM);
    for (uintptr_t i = 0; i < CFG.threads; i++)
        args[i] = (void*)i;

    // actually create the threads
    for (uint32_t j = 1; j < CFG.threads; j++)
        pthread_create(&tid[j], &attr, &run_wrapper, args[j]);

    // all of the other threads should be queued up, waiting to run the
    // benchmark, but they can't until this thread starts the benchmark
    // too...
    run_wrapper(args[0]);

    // Don't call any more transactional functions, because run_wrapper called
    // shutdown.

    // everyone should be done.  Join all threads so we don't leave anything
    // hanging around
    for (uint32_t k = 1; k < CFG.threads; k++)
        pthread_join(tid[k], NULL);

    if (CFG.switch_to_sim) {
        ptlcall_switch_to_native();
    }

    bool v = bench_verify();
    std::cout << "Verification: " << (v ? "Passed" : "Failed") << "\n";

    dump_csv();

    // And call sys shutdown stuff
    TM_SYS_SHUTDOWN();
    return 0;
}

#endif // BMHARNESS_HPP__
